#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
#

"""
嵌套函数代码生成工具方法
"""

import re
from mcpy.ir import IRFunctionNode
from mcpy.ir.expressions.call import IRCall
from mcpy.ir.expressions.basic import IRName
from mcpy.ir.statements.control import IRReturn
from mcpy.codegen.core.context import CodeGenContext
from mcpy.codegen.core.generator import get_global_generator


def get_functor_name(node: IRFunctionNode) -> str:
    """获取稳定的 functor 名称

    优先使用 type_replacement pass 已设置的 functor_name，
    以确保与类型系统中的名称一致。

    Args:
        node: 函数节点

    Returns:
        functor 名称
    """
    # 优先使用已设置的 functor_name（由 type_replacement pass 设置）
    functor_name = node.get_extension_param("functor_name")
    if functor_name:
        return functor_name

    # 查找外层函数，判断是否是全局函数
    parent = node.parent
    while parent and parent.node_type != "function":
        parent = parent.parent

    ir_module = node.get_ir_module()
    module_name = (
        ir_module.module_full_name
        if ir_module and ir_module.module_full_name
        else "module"
    )

    # 如果是全局函数（没有外层函数），使用模块前缀
    if not parent or parent.node_type != "function":
        # 全局函数：{module}_{func_name}
        # 去除特殊字符，替换为下划线
        clean_module_name = module_name.replace(".", "_").replace("/", "_")
        return f"{clean_module_name}_{node.name}"

    # 降级方案：生成嵌套函数名称（应该很少触发，因为 type_replacement 会先运行）
    outer_func_name = parent.name if parent else "func"
    nested_func_name = node.name
    line_num = node.source_line or 0

    # 清理模块名：去除特殊字符，替换为下划线
    clean_module_name = module_name.replace(".", "_").replace("/", "_").replace("-", "_")

    # 生成名称：__nested_func_{module}_{outer}_{nested}_{line}
    return (
        f"__nested_func_{clean_module_name}_{outer_func_name}_{nested_func_name}_{line_num}"
    )


def fix_global_async_struct_call(await_expr, await_expr_code: str, current_node) -> str:
    """修正全局异步函数 struct 的调用语法

    如果 await_expr 是调用全局异步函数，则将 func(args) 改为 struct_name{}(args)

    Args:
        await_expr: await 的表达式（IRCall）
        await_expr_code: 生成的调用代码
        current_node: 当前函数节点，用于符号查找

    Returns:
        修正后的代码
    """
    # 必须是函数调用
    if not isinstance(await_expr, IRCall):
        return await_expr_code

    # 必须是直接函数名调用
    if not isinstance(await_expr.func, IRName):
        return await_expr_code

    func_name = (
        await_expr.func.id
        if isinstance(await_expr.func, IRName)
        else (await_expr.func.name if hasattr(await_expr.func, "name") else "")
    )

    # 检查是否是捕获的变量（嵌套函数）
    # 如果是捕获变量，说明已经是对象实例，不需要修正
    captured_vars = (
        current_node.captured_vars if hasattr(current_node, "captured_vars") else None
    )
    if captured_vars and func_name in captured_vars:
        # 是捕获的嵌套函数，已经是对象实例
        return await_expr_code

    symbol = current_node.lookup_symbol(func_name)
    if not symbol or not symbol.ir_node:
        return await_expr_code

    # 必须是函数节点
    if not isinstance(symbol.ir_node, IRFunctionNode):
        return await_expr_code

    func_node = symbol.ir_node

    # 检查是否是全局函数（没有外层函数）
    parent = func_node.parent
    while parent and parent.node_type != "function":
        parent = parent.parent

    if parent and parent.node_type == "function":
        # 是嵌套函数，检查是否是简单函数（无捕获）
        is_simple = func_node.get_extension_param("is_simple_non_capturing", False)
        if not is_simple:
            # 如果标记未设置，检查 captured_vars 是否为空
            captured_vars = getattr(func_node, "captured_vars", None)
            is_simple = not captured_vars or len(captured_vars) == 0
        
        if is_simple:
            # 简单嵌套函数，生成临时对象调用
            functor_name = get_functor_name(func_node)
            # 将 func_name(args) 改为 functor_name{}(args)
            pattern = rf"^{re.escape(func_name)}\("
            if re.match(pattern, await_expr_code):
                args_part = await_expr_code[len(func_name) :]  # 包含 (args)
                return f"{functor_name}{{}}{args_part}"
            # 如果代码已经是临时对象调用形式，直接返回
            if f"{functor_name}{{}}" in await_expr_code:
                return await_expr_code
        # 非简单嵌套函数，不需要修正（已经是变量或 member）
        return await_expr_code

    # 检查是否是异步函数
    if not func_node.is_async:
        # 不是异步函数
        return await_expr_code

    # 是全局异步函数，需要修正调用语法
    # 获取生成的 struct 名称
    struct_name = get_functor_name(func_node)

    # 将 func_name(args) 改为 struct_name{}(args)
    # 匹配 func_name(args)
    # 注意：需要处理 mcpy::call_async<T>(func_name, args) 的情况
    if await_expr_code.startswith("mcpy::call_async"):
        # 处理 mcpy::call_async<T>(func_name, args) 的情况
        # 将 mcpy::call_async<T>(func_name, args) 改为 mcpy::call_async<T>(struct_name{}, args)
        pattern = rf"mcpy::call_async<[^>]+>\({re.escape(func_name)}"
        if re.search(pattern, await_expr_code):
            await_expr_code = re.sub(
                rf"({re.escape(func_name)})(\(|,|\))",
                rf"{struct_name}{{}}\2",
                await_expr_code,
                count=1,
            )
            return await_expr_code

    # 处理普通函数调用 func_name(args)
    pattern = rf"^{re.escape(func_name)}\("
    if re.match(pattern, await_expr_code):
        args_part = await_expr_code[len(func_name) :]  # 包含 (args)
        return f"{struct_name}{{}}{args_part}"

    return await_expr_code


def indent_code(code: str, level: int) -> str:
    """缩进代码

    Args:
        code: 代码字符串
        level: 缩进级别（每级4空格）

    Returns:
        缩进后的代码
    """
    indent = "    " * level
    lines = code.split("\n")
    indented_lines = [indent + line if line.strip() else "" for line in lines]
    return "\n".join(indented_lines)


def nested_body_generator(node: IRFunctionNode, context: CodeGenContext) -> str:
    """嵌套函数体生成策略：{statements}

    这是嵌套函数与 Lambda 的唯一差异点
    注意：有 await 的 async 函数使用 _generate_async_function_struct 生成，不会到这里

    Args:
        node: 函数节点
        context: 代码生成上下文

    Returns:
        生成的函数体代码
    """
    # 为嵌套函数体创建新的context，设置current_function为嵌套函数名
    nested_context = context.create_child_context(current_function=node.name)

    generator = get_global_generator()

    # 对于无 await 的 async 函数，需要包装返回值为 mc::result
    if node.is_async:
        # 获取返回类型
        ret_type = node.get_cpp_return_type() or "mc::result<void>"
        m = re.match(r"mc::result<\s*(.+)\s*>", ret_type)
        ret_val_type = m.group(1) if m else "void"

        lines = []
        for stmt in node.body.statements:
            if isinstance(stmt, IRReturn) and stmt.expr:
                # 包装返回值
                return_code = generator.generate_expression(stmt.expr, nested_context)
                lines.append(f"return mc::result<{ret_val_type}>({return_code});")
            else:
                stmt_code = generator.generate_statement(stmt, nested_context)
                if stmt_code and not stmt_code.isspace():
                    lines.append(stmt_code)
        return "\n".join(lines)
    else:
        # 非async函数，直接生成
        return generator.generate_statements(node.body.statements, nested_context)
